#include "dag_nodes.h"

#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_statistics.h>
#include <gsl/gsl_linalg.h>

#include <memory>

using namespace dag;
using namespace std;

namespace {
	
	template <typename _type>
	void dealloc_matrix(_type**& matrix, size_t numRows) {
		
		if (matrix != NULL) {
			for (size_t row = 0; row < numRows; row++) {
				delete[] matrix[row];
				matrix[row] = NULL;
			}
		}
		delete[] matrix;
		matrix = NULL;
	}
	
	template<typename _type>
	_type** alloc_matrix(size_t numRows, size_t numCols) {
		_type** matrix = NULL;
		try {
			/* allocate matrix */
			matrix = new _type*[numRows];
			memset(matrix, NULL, numRows * sizeof(_type*));
			for (size_t i = 0; i < numRows; i++) {
				matrix[i] = new _type[numCols];
				memset(matrix[i], 0, numCols * sizeof(_type));
			}
		}
		catch (...) {
			dealloc_matrix<_type>(matrix, numRows);
			throw;
		}
		return matrix;
	}
	
	void center_columns(gsl_matrix* input) {
		/* center columns by subtracting mean */
		for(size_t col = 0; col < input->size2; col++) {
			gsl_vector_view vect = gsl_matrix_column(input, col);
			double mean = gsl_stats_mean(vect.vector.data, 
										 vect.vector.stride, 
										 vect.vector.size);
			int err = gsl_vector_add_constant(&(vect.vector), -mean);
			if(err != 0) {
				/* TODO: better error reporting. Also get error string from gsl
				 */
				throw DagError("gsl math error");
			}
		}
	}
	
	void pca(gsl_matrix* input, gsl_matrix* output) {
		
		if (input->size1 < input->size2 || 
			output->size1 != input->size2 ||
			output->size2 != input->size2 ||
			input->size1 == 0 ||
			input->size2 == 0 ||
			output->size1 == 0 ||
			output->size2 == 0) {
			throw InvalidArgument();
		}
		
		center_columns(input);
		
		gsl_matrix* V = NULL;
		gsl_vector* S = NULL;
		gsl_vector* work = NULL;
		
		try {
			/* alloc working buffers */
			V = gsl_matrix_alloc(input->size2, input->size2);	
			S = gsl_vector_alloc(input->size2);
			work = gsl_vector_alloc(input->size2);
			
			/* perform svd */
			int err = gsl_linalg_SV_decomp(input, V, S, work);
			if (err != 0) {
				/* TODO: better error reporting. Also get error string from gsl
				 */
				throw DagError("gsl math error");
			}
			
			/* copy data into output */
			for (size_t i = 0; i < output->size2; i++) {
				gsl_vector_const_view comp = gsl_matrix_const_column(V, i);
				err = gsl_matrix_set_row(output, i, &comp.vector);
				if (err != 0) {
					/* TODO: better error reporting. Also get error string from gsl
					 */
					throw DagError("gsl math error");
				}
			}
			
			gsl_matrix_free(V);
			V = NULL;
			gsl_vector_free(S);
			S = NULL;
			gsl_vector_free(work);
			work = NULL;
		}
		catch (...) {
			gsl_matrix_free(V);
			V = NULL;
			gsl_vector_free(S);
			S = NULL;
			gsl_vector_free(work);
			work = NULL;			
			throw;
		}
	}
}

PcaNode::PcaNode(const FrameSpec& outSpec, bool forward)
:	m_InputSpec(FrameSpec()),
	m_OutputSpec(FrameSpec()),
	m_Forward(forward),
	m_ForwardOp(NULL),
	m_ReverseOp(NULL)
{
	if (m_Forward) {
		m_Traits.resources = LearnResource_Inputs;
	}
	else {
		m_Traits.resources = LearnResource_Outputs;
	}
	
	m_OutputSpec = outSpec;
	m_Traits.canReverse = true;
	
	if (outSpec.bSize != 0 || outSpec.iSize != 0) {
		throw InvalidArgument();
	}
}

PcaNode::~PcaNode() throw() {
}

const NodeTraits& PcaNode::traits() const {
	return m_Traits;
}

void PcaNode::_init(const FrameSpec& expectedInputSpec) {
	m_InputSpec = expectedInputSpec;
	
	if (m_InputSpec.bSize != 0 || m_InputSpec.iSize != 0) {
		throw InvalidArgument();
	}
	
	if (m_Forward) {
		if (m_OutputSpec.fSize > m_InputSpec.fSize ||
			m_OutputSpec.dSize > m_InputSpec.dSize) {
			throw InvalidArgument();
		}
	}
	else {
		if (m_OutputSpec.fSize < m_InputSpec.fSize ||
			m_OutputSpec.dSize < m_InputSpec.dSize) {
			throw InvalidArgument();
		}
	}
}

const FrameSpec& PcaNode::_inputSpec() const {
	return m_InputSpec;
}

const FrameSpec& PcaNode::_outputSpec() const {
	return m_OutputSpec;
}

auto_ptr<FrameOperator> PcaNode::_createForwardOp() const {
	if (m_ForwardOp.get() == NULL) {
		return auto_ptr<FrameOperator>(NULL);
	}
	return auto_ptr<FrameOperator>(new LinearOp(*m_ForwardOp));
}

auto_ptr<FrameOperator> PcaNode::_createReverseOp() const {
	if (m_ReverseOp.get() == NULL) {
		return auto_ptr<FrameOperator>(NULL);
	}
	return auto_ptr<FrameOperator>(new LinearOp(*m_ReverseOp));
}

void PcaNode::_learn(const FrameDepot* input,  const FrameDepot* output) {
	gsl_matrix* floatInputMat = NULL;
	gsl_matrix* floatTransMat = NULL;
	gsl_matrix* doubleInputMat = NULL;
	gsl_matrix* doubleTransMat = NULL;
	const FrameDepot* source = NULL;
	size_t numFRows = 0;
	size_t numFCols = 0;
	size_t numDRows = 0;
	size_t numDCols = 0;
	float **forFloatTrans = NULL;
	float **revFloatTrans = NULL;
	double **forDoubleTrans = NULL;
	double **revDoubleTrans = NULL;
	
	try {
		/* setup/alloc input and transition matrices based off of whether this
		 * performs pca on the input or the output */
		if (m_Forward) {
			if (m_InputSpec != input->spec()) {
				throw UnequalFrameSpec(m_InputSpec, input->spec());
			}
			source = input;
			if (m_InputSpec.fSize > 0) {
				numFRows = input->size();
				numFCols = m_InputSpec.fSize;
			}
			if (m_InputSpec.dSize > 0) {
				numDRows = input->size();
				numDCols = m_InputSpec.dSize;
			}
		}
		else {
			if (m_OutputSpec != output->spec()) {
				throw UnequalFrameSpec(m_InputSpec, input->spec());
			}
			source = output;
			if (m_OutputSpec.fSize > 0) {
				numFRows = output->size();
				numFCols = m_OutputSpec.fSize;
			}
			if (m_OutputSpec.dSize > 0) {
				numDRows = output->size();
				numDCols = m_OutputSpec.dSize;
			}
		}
		
		/* allocate and fill matrices with data from a frame depot */
		if (source != NULL) {
			if (numFRows > 0 && numFCols > 0) {
				floatInputMat = gsl_matrix_alloc(numFRows, 
												 numFCols);
				floatTransMat = gsl_matrix_alloc(numFCols, 
												 numFCols);
				
				Frame frame(source->spec());
				for (size_t row = 0; row < source->size(); row++) {
					source->fill(row, frame);
					for (size_t col = 0; col < numFCols; col++) {
						gsl_matrix_set(floatInputMat, 
									   row,
									   col, 
									   frame.floats[col]);
					}
				}
			}
			
			if (numDRows > 0 && numDCols > 0) {
				doubleInputMat = gsl_matrix_alloc(numDRows, 
												  numDCols);
				doubleTransMat = gsl_matrix_alloc(numDCols, 
												  numDCols);
				
				Frame frame(source->spec());
				for (size_t row = 0; row < source->size(); row++) {
					source->fill(row, frame);
					for (size_t col = 0; col < numFCols; col++) {
						gsl_matrix_set(doubleInputMat,
									   row,
									   col, 
									   frame.doubles[col]);
					}
				}
			}
		}
		
		/* perform pca on input matrices */
		if (floatInputMat != NULL && floatTransMat != NULL) {
			pca(floatInputMat, floatTransMat);
		}
		
		if (doubleInputMat != NULL && doubleTransMat != NULL) {
			pca(doubleInputMat, doubleTransMat);
		}
		
		/* create matrices and use them to create forward and reverse operators
		 */
		if (floatTransMat != NULL) {
			forFloatTrans = alloc_matrix<float>(m_OutputSpec.fSize,
												m_InputSpec.fSize + 1);
			revFloatTrans = alloc_matrix<float>(m_InputSpec.fSize, 
												m_OutputSpec.fSize + 1);
			if (m_Forward) {
				for (size_t i = 0; i < m_OutputSpec.fSize; i++) {
					for (size_t j = 0; j < m_InputSpec.fSize; j++) {
						forFloatTrans[i][j] = gsl_matrix_get(floatTransMat, i, j);
					}
				}
				
				for (size_t i = 0; i < m_InputSpec.fSize; i++) {
					for (size_t j = 0; j < m_OutputSpec.fSize; j++) {
						revFloatTrans[i][j] = gsl_matrix_get(floatTransMat, j, i);
					}
				}
			}
			else {
				for (size_t i = 0; i < m_OutputSpec.fSize; i++) {
					for (size_t j = 0; j < m_InputSpec.fSize; j++) {
						forFloatTrans[i][j] = gsl_matrix_get(floatTransMat, j, i);
					}
				}
				
				for (size_t i = 0; i < m_InputSpec.fSize; i++) {
					for (size_t j = 0; j < m_OutputSpec.fSize; j++) {
						revFloatTrans[i][j] = gsl_matrix_get(floatTransMat, i, j);
					}
				}
			}
		}
		
		if (doubleTransMat != NULL) {
			forDoubleTrans = alloc_matrix<double>(m_OutputSpec.dSize,
												m_InputSpec.dSize + 1);
			revDoubleTrans = alloc_matrix<double>(m_InputSpec.dSize, 
												m_OutputSpec.dSize + 1);
			if (m_Forward) {
				for (size_t i = 0; i < m_OutputSpec.dSize; i++) {
					for (size_t j = 0; j < m_InputSpec.dSize; j++) {
						forDoubleTrans[i][j] = gsl_matrix_get(doubleTransMat, i, j);
					}
				}
				
				for (size_t i = 0; i < m_InputSpec.dSize; i++) {
					for (size_t j = 0; j < m_OutputSpec.dSize; j++) {
						revDoubleTrans[i][j] = gsl_matrix_get(doubleTransMat, j, i);
					}
				}
			}
			else {
				for (size_t i = 0; i < m_OutputSpec.dSize; i++) {
					for (size_t j = 0; j < m_InputSpec.dSize; j++) {
						forDoubleTrans[i][j] = gsl_matrix_get(doubleTransMat, j, i);
					}
				}
				
				for (size_t i = 0; i < m_InputSpec.dSize; i++) {
					for (size_t j = 0; j < m_OutputSpec.dSize; j++) {
						revDoubleTrans[i][j] = gsl_matrix_get(doubleTransMat, i, j);
					}
				}
			}				
		}
		
		m_ForwardOp.reset(new LinearOp(m_InputSpec, 
									   m_OutputSpec, 
									   forDoubleTrans,
									   forFloatTrans,
									   NULL,
									   NULL));
		m_ReverseOp.reset(new LinearOp(m_OutputSpec,
									   m_InputSpec,
									   revDoubleTrans,
									   revFloatTrans,
									   NULL,
									   NULL));
	
		/* release allocated data */
		dealloc_matrix(forFloatTrans, m_OutputSpec.fSize);
		dealloc_matrix(revFloatTrans, m_InputSpec.fSize);
		dealloc_matrix(forDoubleTrans, m_OutputSpec.dSize);
		dealloc_matrix(revDoubleTrans, m_InputSpec.dSize);
		gsl_matrix_free(floatInputMat);
		floatInputMat = NULL;
		gsl_matrix_free(floatTransMat);
		floatTransMat = NULL;
		gsl_matrix_free(doubleInputMat);
		doubleInputMat = NULL;
		gsl_matrix_free(doubleTransMat);
		doubleTransMat = NULL;
	}
	catch (...) {
		m_ForwardOp.reset(NULL);
		m_ReverseOp.reset(NULL);
		dealloc_matrix(forFloatTrans, m_OutputSpec.fSize);
		dealloc_matrix(revFloatTrans, m_InputSpec.fSize);
		dealloc_matrix(forDoubleTrans, m_OutputSpec.dSize);
		dealloc_matrix(revDoubleTrans, m_InputSpec.dSize);
		gsl_matrix_free(floatInputMat);
		floatInputMat = NULL;
		gsl_matrix_free(floatTransMat);
		floatTransMat = NULL;
		gsl_matrix_free(doubleInputMat);
		doubleInputMat = NULL;
		gsl_matrix_free(doubleTransMat);
		doubleTransMat = NULL;
		
		throw;
	}
}